#include <Arduino.h>
#include "encoder.h"
#include "tim.h"
#include "adc.h"
#include "gpio.h"  // Add GPIO header
#include "stm32g4xx_hal.h"
#include "sin_table.h"  // Add sine table header
#include "lowpass.h"  // Add include for low-pass filter
#include "PI.h"
#include "cordic.h"  // Add include

// Add before other includes:
extern "C" {
    #include "foc.h"         // Include MATLAB generated header
}


#define CTRL_FREQ 10000
#define Ubase 6.92

#define CH_COUNT 5  // Three channels: angle, debug, state
#define CMD_LENGTH 2  // Command length in bytes (16 bits)

// Command types and states
#define CMD_STATE 0x01    // State machine command
#define CMD_UD    0x02    // u_d control command
#define CMD_IQ    0x03    // i_q control command

enum State {
    STANDBY = 0,
    CALIBRE = 1,
    OPENLOOP_RUN = 2,  // Renamed from RUNNING
    CLOSELOOP_RUN = 3,
    SHUTDOWN = 4
};

struct Frame {
    float fdata[CH_COUNT];
    unsigned char tail[4]{0x00, 0x00, 0x80, 0x7f};
};

Encoder encoder;
Frame dataFrame;
State currentState = STANDBY;
uint8_t rxBuffer[CMD_LENGTH];  // Buffer for 2 bytes command
uint8_t rxCount = 0;

// Global variables for sawtooth wave and sine calculation
volatile float theta_self = 0.0f;
volatile float sin_result = 0.0f;  // Add volatile result variable
volatile float sin_cordic_result = 0.0f;  // Add volatile result variable
volatile float cos_cordic_result = 0.0f;  // Add volatile result variable
const float delta_theta = 2.0f * PI * (1.0f/500.0f);  // Angular step per PWM cycle

// Add encoder offset angle
volatile float theta_offset = 0.0f;

// Add control angle
volatile float theta_ctrl = 0.0f;

// adc
float sample_gain = 3.3/4095;
float sample_vdc_gain = (180.0 + 12.0) / 12.0; 
float sample_i_gain = 1.0/(0.33*2*(2200.0/(2200.0+680.0)));
float sample_i_bias = 3.3*680.0/(680.0+2200.0)*2/3.3*4095;


// Control parameters
const float VDC0 = 12.0f;                    // Default DC bus voltage
const float TWO_PI_DIV_3 = 2.0f * PI / 3.0f; // 2π/3 constant
volatile float f_olp = 56.0f;                  // Open loop frequency in Hz
volatile float k_omega_e_olp = 0.0f;             // Speed reference coefficient
volatile float omega_e_olp = 0.0f;               // Actual speed reference
volatile float u_d = 0.0f, i_q = 0.0f;       // DQ voltage commands
volatile float u_a = 0.0f, u_b = 0.0f, u_c = 0.0f; // ABC voltage commands
volatile float d_a = 0.0f, d_b = 0.0f, d_c = 0.0f; // Duty cycles

volatile float i_a = 0.0f, i_b = 0.0f, i_c = 0.0f; // Phase currents
float i_q_fdb = 0.0f; // Phase current in q-axis

volatile bool pwm_enabled = false;
// Update i_q based on control mode
float i_q_mod;

// Add variables for speed calculation and filtering
volatile float omega_e = 0.0f;     // Electrical speed in rad/s
volatile float omega_ef = 0.0f;    // 2Hz Filtered electrical speed 
volatile float omega_ef2 = 0.0f;    // 10Hz Filtered electrical speed
volatile float last_theta_e = 0.0f; // Previous electrical angle
LowPassFilter speedFilter;         // 2Hz low-pass filter (uses default constructor)
LowPassFilter speedFilter2(50.0f); // 10Hz low-pass filter

// Add LED control variables
volatile uint32_t led_counter = 0;
const uint32_t LED_TOGGLE_COUNT = 5000; // 5kHz * 1sec = 5000 counts
volatile bool led_state = false;

// Update PI controller with kp=0.1, ki=0.01/5000, limit=±5
PIController speedController(0.00001f, 0.001f/CTRL_FREQ, 0.5f);
float eer_lim = 10.0f;

// Add reference speed variable
volatile float omega_er = 0.0f;

float lookup_sin(float angle) {
    // Normalize angle to 0-360 degrees first
    float deg = angle * 180.0f / PI;
    while(deg >= 360.0f) deg -= 360.0f;
    while(deg < 0.0f) deg += 360.0f;
    
    // Calculate table index (91 points total)
    int idx = (int)(deg / 4.0f);
    // Ensure index is within bounds
    if(idx >= 91) idx = 90;  // Clamp to last valid index
    if(idx < 0) idx = 0;     // Clamp to first valid index
    
    return sin_table[idx];
}

void processCommand() {
    uint8_t cmdType = rxBuffer[0];    
    uint8_t cmdData = rxBuffer[1];    
    
    switch(cmdType) {
        case CMD_STATE:
            switch(cmdData) {
                case 0x00:  // Enter standby
                    if(pwm_enabled) {
                        // Disable PWM outputs
                        HAL_TIM_PWM_Stop(&htim1, TIM_CHANNEL_1);
                        HAL_TIM_PWM_Stop(&htim1, TIM_CHANNEL_2);
                        HAL_TIM_PWM_Stop(&htim1, TIM_CHANNEL_3);
                        __HAL_TIM_MOE_DISABLE(&htim1);
                        
                        // Disable PWM enable pins and PB5
                        HAL_GPIO_WritePin(GPIOB, GPIO_PIN_13|GPIO_PIN_14|GPIO_PIN_15|GPIO_PIN_5, GPIO_PIN_RESET);
                        
                        pwm_enabled = false;
                    }
                    currentState = STANDBY;
                    theta_self = 0.0f;
                    theta_ctrl = theta_self;  // Use self angle in standby
                    theta_offset = encoder.getDefaultOffset();  // Always use default offset
                    speedController.reset();
                    break;
                    
                case 0x01:  // Enter calibration
                    if(!pwm_enabled) {
                        // Enable PWM enable pins and PB5
                        HAL_GPIO_WritePin(GPIOB, GPIO_PIN_13|GPIO_PIN_14|GPIO_PIN_15|GPIO_PIN_5, GPIO_PIN_SET);
                        
                        // Enable PWM outputs
                        HAL_TIM_PWM_Start(&htim1, TIM_CHANNEL_1);
                        HAL_TIM_PWM_Start(&htim1, TIM_CHANNEL_2);
                        HAL_TIM_PWM_Start(&htim1, TIM_CHANNEL_3);
                        __HAL_TIM_MOE_ENABLE(&htim1);
                        
                        pwm_enabled = true;
                    }
                    currentState = CALIBRE;
                    theta_self = 0.0f;
                    theta_ctrl = theta_self;  // Use self angle in calibration
                    // Note: theta_offset will be updated in interrupt during calibration
                    speedController.reset();
                    break;
                    
                case 0x02:  // Enter openloop run
                    if(!pwm_enabled) {
                        // Enable PWM if not enabled
                        HAL_GPIO_WritePin(GPIOB, GPIO_PIN_13|GPIO_PIN_14|GPIO_PIN_15|GPIO_PIN_5, GPIO_PIN_SET);
                        HAL_TIM_PWM_Start(&htim1, TIM_CHANNEL_1);
                        HAL_TIM_PWM_Start(&htim1, TIM_CHANNEL_2);
                        HAL_TIM_PWM_Start(&htim1, TIM_CHANNEL_3);
                        __HAL_TIM_MOE_ENABLE(&htim1);
                        pwm_enabled = true;
                    }
                    
                    // Check if we're coming from calibration before changing state
                    if(currentState != CALIBRE) {
                        theta_offset = encoder.getDefaultOffset();
                    }
                    
                    currentState = OPENLOOP_RUN;
                    u_d = 0.0f;  // Reset u_d to 0 when entering running state
                    theta_ctrl = encoder.getElectricalAngle() - theta_offset;  // Already in radians
                    break;

                case 0x03:  // Enter closeloop run 
                    if(!pwm_enabled) {
                        // Enable PWM if not enabled
                        HAL_GPIO_WritePin(GPIOB, GPIO_PIN_13|GPIO_PIN_14|GPIO_PIN_15|GPIO_PIN_5, GPIO_PIN_SET);
                        HAL_TIM_PWM_Start(&htim1, TIM_CHANNEL_1);
                        HAL_TIM_PWM_Start(&htim1, TIM_CHANNEL_2);
                        HAL_TIM_PWM_Start(&htim1, TIM_CHANNEL_3);
                        __HAL_TIM_MOE_ENABLE(&htim1);
                        pwm_enabled = true;
                    }
                    
                    if(currentState != CALIBRE) {
                        theta_offset = encoder.getDefaultOffset();
                    }
                    
                    currentState = CLOSELOOP_RUN;
                    u_d = 0.0f;  // Reset u_d
                    speedController.reset();  // Reset speed controller
                    theta_ctrl = encoder.getElectricalAngle() - theta_offset;
                    break;
                    
                default:
                    currentState = STANDBY;
                    speedController.reset();
                    theta_ctrl = theta_self;  // Use self angle in standby
                    theta_offset = encoder.getDefaultOffset();  // Use default offset
                    break;
            }
            break;
            
        case CMD_UD:  // 0x02: u_d control command
            switch(cmdData) {
                case 0x01: u_d = 0.1f; break;
                case 0x02: u_d = 0.2f; break;
                case 0x03: u_d = 0.3f; break;
                case 0x04: u_d = 0.4f; break;
                case 0x05: u_d = 0.5f; break;
                default: u_d = 0.0f; break;
            }
            break;

        case CMD_IQ:  // 0x03: i_q control command
            switch(cmdData) {
                case 0x01: i_q = 0.1f; break;
                case 0x02: i_q = 0.2f; break;
                case 0x03: i_q = 0.3f; break;
                case 0x04: i_q = 0.4f; break;
                case 0x05: i_q = 0.5f; break;
                default: i_q = 0.0f; break;
            }
            break;
    }
    rxCount = 0;
}

void setup() {
    Serial.begin(2000000);  // Using default UART
    encoder.init();
    
    GPIO_PWM_EN_Init();     // Initialize PWM enable pins
    
    // Set PWM enable pins high and keep them high
    HAL_GPIO_WritePin(GPIOB, GPIO_PIN_13, GPIO_PIN_SET);
    HAL_GPIO_WritePin(GPIOB, GPIO_PIN_14, GPIO_PIN_SET);
    HAL_GPIO_WritePin(GPIOB, GPIO_PIN_15, GPIO_PIN_SET);
    
    ADC1_Init();           // Initialize ADC1
    TIM1_PWM_Init();       // Initialize TIM1 PWM
    TIM1_Start();          // Start PWM output

    // Configure PA5 as output for LED
    GPIO_InitTypeDef GPIO_InitStruct = {0};
    __HAL_RCC_GPIOA_CLK_ENABLE();
    GPIO_InitStruct.Pin = GPIO_PIN_5;
    GPIO_InitStruct.Mode = GPIO_MODE_OUTPUT_PP;
    GPIO_InitStruct.Pull = GPIO_NOPULL;
    GPIO_InitStruct.Speed = GPIO_SPEED_FREQ_LOW;
    HAL_GPIO_Init(GPIOA, &GPIO_InitStruct);

    CORDIC_Init();  // Initialize CORDIC
    foc_initialize();
}

// Add timestamp variable for UART timing
uint32_t last_uart_time = 0;
const uint32_t UART_INTERVAL = 500; // 500us interval

void loop() {
    // Process incoming serial data
    while (Serial.available() >= CMD_LENGTH) {
        // Read complete command at once
        rxCount = Serial.readBytes(rxBuffer, CMD_LENGTH);
        
        if (rxCount == CMD_LENGTH) {
            processCommand();
        }
    }

    // Read both ADC values
    ADC1_ReadValues();

    // Check if 2ms has elapsed
    uint32_t current_time = micros();
    if(current_time - last_uart_time >= UART_INTERVAL) {
        // Update frame data
        dataFrame.fdata[0] = omega_er;                          
        dataFrame.fdata[1] = omega_ef2; // DC bus voltage
        dataFrame.fdata[2] = i_q_mod;                            // Send actual speed reference
        dataFrame.fdata[3] = i_q_fdb;                           // Send filtered electrical speed
        dataFrame.fdata[4] = i_c;  
        // Send frame data
        Serial.write((uint8_t*)&dataFrame, sizeof(Frame));
        
        last_uart_time = current_time;
    }
    // No delay needed - timing is handled by timestamp check
}

extern "C" void ADC1_2_IRQHandler(void)
{
    if(ADC1->ISR & ADC_ISR_JEOS) 
    {
        ExtU local_rtU;  // Create local instance of input structure
        ExtY local_rtY;  // Create local instance of output structure
        local_rtU.enable = 0.0;
        // LED control using interrupt counter
        led_counter++;
        if(led_counter >= LED_TOGGLE_COUNT) {
            led_counter = 0;
            led_state = !led_state;
            HAL_GPIO_WritePin(GPIOA, GPIO_PIN_5, led_state ? GPIO_PIN_SET : GPIO_PIN_RESET);
        }

        iu_raw = HAL_ADCEx_InjectedGetValue(&hadc1, ADC_INJECTED_RANK_1);
        iv_raw = HAL_ADCEx_InjectedGetValue(&hadc1, ADC_INJECTED_RANK_2);
        iw_raw = HAL_ADCEx_InjectedGetValue(&hadc1, ADC_INJECTED_RANK_3);


        i_a = (iu_raw-sample_i_bias)*sample_gain*sample_i_gain-0.060;
        i_b = (iv_raw-sample_i_bias)*sample_gain*sample_i_gain-0.055;
        i_c = (iw_raw-sample_i_bias)*sample_gain*sample_i_gain-0.023;
        
        // Update speed reference
        k_omega_e_olp = speedref_raw / 4095.0f;
        omega_e_olp = k_omega_e_olp * f_olp *2.0f * PI;

        // Update theta for sawtooth wave based on omega_e_olp
        if(currentState != CALIBRE) {
            theta_self += omega_e_olp / CTRL_FREQ;  //10kHz PWM frequency
            if(theta_self >= 2.0f * PI) {
                theta_self -= 2.0f * PI;
            }
        }
        encoder.updateAngles();

        // Calculate electrical angular speed
        float current_theta_e = encoder.getElectricalAngle();
        float delta_theta_e = current_theta_e - last_theta_e;
        
        // Normalize delta_theta to -π to π
        if(delta_theta_e > PI) delta_theta_e -= 2.0f * PI;
        if(delta_theta_e < -PI) delta_theta_e += 2.0f * PI;
        
        // Calculate speed in rad/s (10kHz sampling)
        omega_e = delta_theta_e * CTRL_FREQ;
        
        // Apply low-pass filter
        omega_ef = speedFilter.update(omega_e);
        
        // Store current angle for next calculation
        last_theta_e = current_theta_e;

        // Only update offset during calibration
        if(currentState == CALIBRE) {
            theta_offset = encoder.getElectricalAngle(); // Update in radians
        }
        
        // Always use current theta_offset (whether from calibration or default)
        if(currentState == OPENLOOP_RUN || currentState == CLOSELOOP_RUN) {
            theta_ctrl = encoder.getElectricalAngle() - theta_offset;  // Already in radians
            local_rtU.enable = 1.0;
        } else {
            theta_ctrl = theta_self;
            local_rtU.enable = 0.0;
        }

        // Calculate speed reference for closed loop control
        omega_er = (speedref_raw / 4095.0f) * 7.0f * 2.0f * PI * 10.0f;
        
        // Apply second low-pass filter for PI control
        omega_ef2 = speedFilter2.update(omega_e);

        // Interface with foc


        if(currentState == CLOSELOOP_RUN) {
            // Limit speed error to ±2
            float speed_error = (omega_er - omega_ef2)/10;
            if(speed_error > eer_lim) speed_error = eer_lim;
            if(speed_error < -eer_lim) speed_error = -eer_lim;
            
            // Use limited error for PI control
            i_q_mod = speedController.update(omega_ef2 + speed_error, omega_ef2);            
        } else {
            i_q_mod = i_q * speedref_raw * sample_gain / 3.3f;  // Original scaling in open loop
        }



        // Use CORDIC to calculate sine
        CORDIC_CalcCosSin(theta_ctrl,&cos_cordic_result,&sin_cordic_result);


        
        local_rtU.ia = -i_a;  
        local_rtU.ib = -i_b;
        local_rtU.ic = -i_c;
        local_rtU.cosine = cos_cordic_result;
        local_rtU.sine = sin_cordic_result;
        local_rtU.idcRef = i_q_mod;

        // Copy local inputs to global structure
        rtU = local_rtU;
        
        // Execute ZERO crossing detection
        foc_step();

        // Copy outputs if needed
        local_rtY = rtY;
        i_q_fdb = local_rtY.idc;  // Update i_q from output structure
        d_a = local_rtY.S[0];  // Get d_a from output structure
        d_b = local_rtY.S[2];  // Get d_b from output structure
        d_c = local_rtY.S[4];  // Get d_b from output structure
        

        // Update TIM1 compare values
        uint16_t pwm_period = TIM1->ARR;
        TIM1->CCR1 = (uint16_t)(d_a * pwm_period);
        TIM1->CCR2 = (uint16_t)(d_b * pwm_period);
        TIM1->CCR3 = (uint16_t)(d_c * pwm_period);


        //sin_result = lookup_sin(theta_ctrl);

        ADC1->ISR = ADC_ISR_JEOS;
    }
}
